kreuzberg 2.1.2__py3-none-any.whl → 3.0.1__py3-none-any.whl

kreuzberg/__init__.py

@@ -1,5 +1,10 @@
-from ._tesseract import PSMMode
-from ._types import ExtractionResult, Metadata
+from kreuzberg._ocr._easyocr import EasyOCRConfig
+from kreuzberg._ocr._paddleocr import PaddleOCRConfig
+from kreuzberg._ocr._tesseract import TesseractConfig
+
+from ._ocr._tesseract import PSMMode
+from ._registry import ExtractorRegistry
+from ._types import ExtractionConfig, ExtractionResult, Metadata
 from .exceptions import KreuzbergError, MissingDependencyError, OCRError, ParsingError, ValidationError
 from .extraction import (
     batch_extract_bytes,
@@ -7,22 +12,31 @@ from .extraction import (
     batch_extract_file,
     batch_extract_file_sync,
     extract_bytes,
+    extract_bytes_sync,
     extract_file,
+    extract_file_sync,
 )
 
 __all__ = [
+    "EasyOCRConfig",
+    "ExtractionConfig",
     "ExtractionResult",
+    "ExtractorRegistry",
     "KreuzbergError",
     "Metadata",
     "MissingDependencyError",
     "OCRError",
     "PSMMode",
+    "PaddleOCRConfig",
     "ParsingError",
+    "TesseractConfig",
     "ValidationError",
     "batch_extract_bytes",
     "batch_extract_bytes_sync",
     "batch_extract_file",
     "batch_extract_file_sync",
     "extract_bytes",
+    "extract_bytes_sync",
     "extract_file",
+    "extract_file_sync",
 ]

kreuzberg/_chunker.py

@@ -0,0 +1,51 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from kreuzberg import MissingDependencyError
+from kreuzberg._constants import DEFAULT_MAX_CHARACTERS, DEFAULT_MAX_OVERLAP
+from kreuzberg._mime_types import MARKDOWN_MIME_TYPE
+
+if TYPE_CHECKING:
+    from semantic_text_splitter import MarkdownSplitter, TextSplitter
+
+_chunkers: dict[tuple[int, int, str], MarkdownSplitter | TextSplitter] = {}
+
+
+def get_chunker(
+    mime_type: str,
+    max_characters: int = DEFAULT_MAX_CHARACTERS,
+    overlap_characters: int = DEFAULT_MAX_OVERLAP,
+) -> MarkdownSplitter | TextSplitter:
+    """Creates and returns a Chunker object configured with the given maximum
+    characters per chunk and overlap between chunks.
+
+    Args:
+        mime_type: The mime type of the content.
+        max_characters: Maximum number of characters allowed in each chunk.
+        overlap_characters: Number of characters overlapping between two consecutive chunks.
+
+    Raises:
+        MissingDependencyError: if semantic-text-splitter is not installed.
+
+    Returns:
+        Chunker: A Chunker object configured with the specified maximum
+            characters and overlap.
+    """
+    key = (max_characters, overlap_characters, mime_type)
+    if key not in _chunkers:
+        try:
+            if mime_type == MARKDOWN_MIME_TYPE:
+                from semantic_text_splitter import MarkdownSplitter
+
+                _chunkers[key] = MarkdownSplitter(max_characters, overlap_characters)
+            else:
+                from semantic_text_splitter import TextSplitter
+
+                _chunkers[key] = TextSplitter(max_characters, overlap_characters)
+        except ImportError as e:
+            raise MissingDependencyError.create_for_package(
+                dependency_group="chunking", functionality="chunking", package_name="semantic-text-splitter"
+            ) from e
+
+    return _chunkers[key]

kreuzberg/_constants.py

@@ -1,8 +1,7 @@
 from __future__ import annotations
 
-from multiprocessing import cpu_count
 from typing import Final
 
-DEFAULT_MAX_PROCESSES: Final[int] = cpu_count()
-MINIMAL_SUPPORTED_TESSERACT_VERSION: Final[int] = 5
 MINIMAL_SUPPORTED_PANDOC_VERSION: Final[int] = 2
+DEFAULT_MAX_CHARACTERS: Final[int] = 2000
+DEFAULT_MAX_OVERLAP: Final[int] = 100

kreuzberg/_extractors/_base.py

@@ -0,0 +1,92 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, ClassVar
+
+if TYPE_CHECKING:
+    from pathlib import Path
+
+    from kreuzberg import ExtractionResult
+    from kreuzberg._types import ExtractionConfig
+
+
+class Extractor(ABC):
+    """Abstract base class for file content extraction.
+
+    This class provides the interface for different types of content extractors.
+    Subclasses are expected to implement the methods for extracting content
+    either asynchronously or synchronously and determining the supported MIME types.
+
+    Attributes:
+        SUPPORTED_MIME_TYPES: The set of supported mime types - all none abstract extractors must implement this.
+
+    Args:
+        mime_type: The MIME type that this extractor handles (e.g., "application/pdf").
+        config: Configuration options for the extraction process.
+    """
+
+    __slots__ = ("config", "mime_type")
+
+    SUPPORTED_MIME_TYPES: ClassVar[set[str]]
+
+    def __init__(self, mime_type: str, config: ExtractionConfig) -> None:
+        self.mime_type = mime_type
+        self.config = config
+
+    @abstractmethod
+    async def extract_bytes_async(self, content: bytes) -> ExtractionResult:
+        """Asynchronously extract content from a byte stream.
+
+        Args:
+            content: The byte content to extract.
+
+        Returns:
+            ExtractionResult: The extracted content along with metadata about the extraction.
+        """
+
+    @abstractmethod
+    async def extract_path_async(self, path: Path) -> ExtractionResult:
+        """Asynchronously extract content from a file located at the specified path.
+
+        Args:
+            path: The path to the file to process.
+
+        Returns:
+            ExtractionResult: The extracted content along with metadata about the extraction.
+        """
+
+    @abstractmethod
+    def extract_bytes_sync(self, content: bytes) -> ExtractionResult:
+        """Synchronously extract content from a byte stream.
+
+        Args:
+            content: The byte content to extract.
+
+        Returns:
+            ExtractionResult: The extracted content along with metadata about the extraction.
+        """
+
+    @abstractmethod
+    def extract_path_sync(self, path: Path) -> ExtractionResult:
+        """Synchronously extract content from a file located at the specified path.
+
+        Args:
+            path: The path to the file to process.
+
+        Returns:
+            ExtractionResult: The extracted content along with metadata about the extraction.
+        """
+
+    @classmethod
+    def supports_mimetype(cls, mime_type: str) -> bool:
+        """Verify whether the extractor supports the given MIME type.
+
+        Args:
+            mime_type: The MIME type to check (e.g., "application/pdf").
+
+        Returns:
+            bool: True if the MIME type is supported, False otherwise.
+        """
+        return mime_type in cls.SUPPORTED_MIME_TYPES or any(
+            mime_type.startswith(supported_type) for supported_type in cls.SUPPORTED_MIME_TYPES
+        )

kreuzberg/_extractors/_html.py

@@ -0,0 +1,34 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, ClassVar
+
+import html_to_markdown
+from anyio import Path as AsyncPath
+
+from kreuzberg._extractors._base import Extractor
+from kreuzberg._mime_types import HTML_MIME_TYPE, MARKDOWN_MIME_TYPE
+from kreuzberg._types import ExtractionResult
+from kreuzberg._utils._string import normalize_spaces, safe_decode
+from kreuzberg._utils._sync import run_sync
+
+if TYPE_CHECKING:
+    from pathlib import Path
+
+
+class HTMLExtractor(Extractor):
+    SUPPORTED_MIME_TYPES: ClassVar[set[str]] = {HTML_MIME_TYPE}
+
+    async def extract_bytes_async(self, content: bytes) -> ExtractionResult:
+        return await run_sync(self.extract_bytes_sync, content)
+
+    async def extract_path_async(self, path: Path) -> ExtractionResult:
+        content = await AsyncPath(path).read_bytes()
+        return await run_sync(self.extract_bytes_sync, content)
+
+    def extract_bytes_sync(self, content: bytes) -> ExtractionResult:
+        result = html_to_markdown.convert_to_markdown(safe_decode(content))
+        return ExtractionResult(content=normalize_spaces(result), mime_type=MARKDOWN_MIME_TYPE, metadata={}, chunks=[])
+
+    def extract_path_sync(self, path: Path) -> ExtractionResult:
+        content = path.read_bytes()
+        return self.extract_bytes_sync(content)

kreuzberg/_extractors/_image.py

@@ -0,0 +1,74 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, ClassVar
+
+import anyio
+from anyio import Path as AsyncPath
+
+from kreuzberg._extractors._base import Extractor
+from kreuzberg._mime_types import IMAGE_MIME_TYPES
+from kreuzberg._ocr import get_ocr_backend
+from kreuzberg._utils._tmp import create_temp_file
+from kreuzberg.exceptions import ValidationError
+
+if TYPE_CHECKING:  # pragma: no cover
+    from collections.abc import Mapping
+    from pathlib import Path
+
+    from kreuzberg._types import ExtractionResult
+
+
+class ImageExtractor(Extractor):
+    SUPPORTED_MIME_TYPES: ClassVar[set[str]] = IMAGE_MIME_TYPES
+
+    IMAGE_MIME_TYPE_EXT_MAP: ClassVar[Mapping[str, str]] = {
+        "image/bmp": "bmp",
+        "image/x-bmp": "bmp",
+        "image/x-ms-bmp": "bmp",
+        "image/gif": "gif",
+        "image/jpeg": "jpg",
+        "image/pjpeg": "jpg",
+        "image/png": "png",
+        "image/tiff": "tiff",
+        "image/x-tiff": "tiff",
+        "image/jp2": "jp2",
+        "image/jpx": "jpx",
+        "image/jpm": "jpm",
+        "image/mj2": "mj2",
+        "image/webp": "webp",
+        "image/x-portable-anymap": "pnm",
+        "image/x-portable-bitmap": "pbm",
+        "image/x-portable-graymap": "pgm",
+        "image/x-portable-pixmap": "ppm",
+    }
+
+    async def extract_bytes_async(self, content: bytes) -> ExtractionResult:
+        extension = self._get_extension_from_mime_type(self.mime_type)
+        file_path, unlink = await create_temp_file(f".{extension}")
+        await AsyncPath(file_path).write_bytes(content)
+        try:
+            return await self.extract_path_async(file_path)
+        finally:
+            await unlink()
+
+    async def extract_path_async(self, path: Path) -> ExtractionResult:
+        if self.config.ocr_backend is None:
+            raise ValidationError("ocr_backend is None, cannot perform OCR")
+
+        return await get_ocr_backend(self.config.ocr_backend).process_file(path, **self.config.get_config_dict())
+
+    def extract_bytes_sync(self, content: bytes) -> ExtractionResult:
+        return anyio.run(self.extract_bytes_async, content)
+
+    def extract_path_sync(self, path: Path) -> ExtractionResult:
+        return anyio.run(self.extract_path_async, path)
+
+    def _get_extension_from_mime_type(self, mime_type: str) -> str:
+        if mime_type in self.IMAGE_MIME_TYPE_EXT_MAP:
+            return self.IMAGE_MIME_TYPE_EXT_MAP[mime_type]
+
+        for k, v in self.IMAGE_MIME_TYPE_EXT_MAP.items():
+            if k.startswith(mime_type):
+                return v
+
+        raise ValidationError("unsupported mimetype", context={"mime_type": mime_type})